# -*- coding: utf-8 -*-
"""
阿里云语音识别WebSocket服务器
基于阿里云Paraformer实时语音识别SDK
"""

import base64
import json
import asyncio
import uuid
import logging
from fastapi import WebSocket, WebSocketDisconnect, APIRouter
from .aliyun_asr_client import asr_manager, AliyunASRClient

logger = logging.getLogger(__name__)

# 创建路由器
aliyun_ws_asr_router = APIRouter()

@aliyun_ws_asr_router.websocket("/ws/aliyun_asr", name="阿里云语音识别WebSocket接口")
async def websocket_endpoint(websocket: WebSocket):
    """
    阿里云语音识别WebSocket接口
    
    客户端发送JSON格式消息：
    {
        "type": "heartbeat"  // 心跳包
    }
    或
    {
        "type": "audio",
        "frame_type": 0|1|2,  // 0x00 起始帧, 0x01 中间帧, 0x02 结束帧
        "data": "<base64_encoded_audio_chunk>",
        "config": {           // 可选配置
            "model": "paraformer-realtime-v2",
            "format": "pcm",
            "sample_rate": 16000,
            "language_hints": ["zh", "en"]
        }
    }
    """
    await websocket.accept()
    client_id = str(uuid.uuid4())
    logger.info(f"[阿里云ASR] WebSocket连接已接受，客户端ID: {client_id}")
    
    buffer = bytearray()
    client = None
    last_sent_text = ""  # 记录上次发送的文本，避免重复发送
    config = {}  # 客户端配置

    # 定义结果回调函数
    async def send_result_to_client(text: str, sentence_end: bool = False):
        logger.info(f"回调收到: text={text}, sentence_end={sentence_end}")
        nonlocal last_sent_text
        try:
            # 避免重复发送相同的文本（仅在sentence_end为False时）
            if text != last_sent_text or sentence_end:
                message = json.dumps({
                    "type": "result",
                    "text": text,
                    "sentence_end": sentence_end,
                    "client_id": client_id
                })
                logger.info(f"[阿里云ASR] 发送给客户端 {client_id}: {text} (sentence_end={sentence_end})")
                await websocket.send_text(message)
                last_sent_text = text if not sentence_end else ""
            else:
                logger.debug(f"[阿里云ASR] 跳过重复文本: '{text}' (sentence_end={sentence_end})")
        except Exception as e:
            logger.error(f"[阿里云ASR] 发送结果失败: {e}")

    try:
        while True:
            raw_msg = await websocket.receive_text()
            try:
                msg_obj = json.loads(raw_msg)
            except json.JSONDecodeError as e:
                logger.error(f"[阿里云ASR] JSON解析失败: {e}，跳过该消息")
                continue

            msg_type = msg_obj.get("type")
            
            if msg_type == "heartbeat":
                # 心跳包响应
                await websocket.send_text(json.dumps({
                    "type": "pong",
                    "client_id": client_id
                }))
                continue
                
            elif msg_type == "audio":
                frame_type = msg_obj.get("frame_type")
                b64_data = msg_obj.get("data", "")
                
                # 处理配置信息
                if "config" in msg_obj:
                    config.update(msg_obj["config"])
                
                try:
                    audio_bytes = base64.b64decode(b64_data)
                except Exception as e:
                    logger.error(f"[阿里云ASR] base64解码失败: {e}, 跳过此帧")
                    continue
                
                if frame_type == 0:  # 起始帧
                    buffer.clear()
                    last_sent_text = ""  # 重置上次发送的文本
                    
                    # 创建新的ASR客户端
                    try:
                        client = await asr_manager.create_client(
                            client_id,
                            model=config.get("model", "paraformer-realtime-v2"),
                            format=config.get("format", "pcm"),
                            sample_rate=config.get("sample_rate", 16000),
                            language_hints=config.get("language_hints", ["zh", "en"]),
                            on_result_callback=send_result_to_client
                        )
                        await client.connect()
                        
                        # 开始识别会话，参考官方demo
                        start_success = await client.start()
                        if not start_success:
                            raise Exception("识别会话启动失败")
                        
                        # 发送连接成功消息
                        await websocket.send_text(json.dumps({
                            "type": "connected",
                            "client_id": client_id,
                            "model": client.model,
                            "sample_rate": client.sample_rate
                        }))
                        
                        logger.info(f"[阿里云ASR] 客户端 {client_id} 连接成功并开始识别")
                        
                    except Exception as e:
                        logger.error(f"[阿里云ASR] 客户端 {client_id} 连接失败: {e}")
                        await websocket.send_text(json.dumps({
                            "type": "error",
                            "message": f"连接失败: {str(e)}",
                            "client_id": client_id
                        }))
                        continue
                
                # 处理音频数据
                if client is not None:
                    buffer.extend(audio_bytes)
                    
                    # 按200ms分包，根据采样率计算包大小
                    chunk_size = 1600 if config.get("sample_rate", 16000) == 8000 else 3200
                    
                    while len(buffer) >= chunk_size:
                        chunk = bytes(buffer[:chunk_size])
                        await client.send_audio(chunk)
                        buffer = buffer[chunk_size:]
                    
                    if frame_type == 2:  # 结束帧
                        if len(buffer) > 0:
                            await client.send_audio(bytes(buffer), is_last=True)
                            buffer.clear()
                        
                        # 关闭客户端
                        await asr_manager.remove_client(client_id)
                        client = None
                        
                        # 发送识别完成消息
                        await websocket.send_text(json.dumps({
                            "type": "completed",
                            "client_id": client_id
                        }))
                        
                        logger.info(f"[阿里云ASR] 客户端 {client_id} 识别完成")
                        continue
                        
            elif msg_type == "config":
                # 更新配置
                config.update(msg_obj.get("config", {}))
                await websocket.send_text(json.dumps({
                    "type": "config_updated",
                    "client_id": client_id,
                    "config": config
                }))
                
            else:
                logger.warning(f"[阿里云ASR] 未知消息类型: {msg_type}, 跳过处理")
                continue
                
    except WebSocketDisconnect:
        logger.info(f"[阿里云ASR] 客户端 {client_id} WebSocket已断开")
        if client_id in asr_manager.clients:
            await asr_manager.remove_client(client_id)
            
    except Exception as e:
        logger.error(f"[阿里云ASR] 处理WebSocket消息时发生错误: {e}")
        if client_id in asr_manager.clients:
            await asr_manager.remove_client(client_id)


@aliyun_ws_asr_router.websocket("/ws/aliyun_asr_v2", name="阿里云语音识别WebSocket接口_v2")
async def websocket_endpoint_v2(websocket: WebSocket):
    """
    阿里云语音识别WebSocket接口 v2 - 简化版本
    
    客户端发送JSON格式消息：
    {
        "type": "start",     // 开始识别
        "config": {          // 配置参数
            "model": "paraformer-realtime-v2",
            "format": "pcm",
            "sample_rate": 16000,
            "language_hints": ["zh", "en"]
        }
    }
    
    {
        "type": "audio",     // 音频数据
        "data": "<base64_encoded_audio_chunk>"
    }
    
    {
        "type": "stop"       // 停止识别
    }
    """
    logger.info("[阿里云ASR v2] 收到WebSocket连接请求")
    await websocket.accept()
    client_id = str(uuid.uuid4())
    logger.info(f"[阿里云ASR v2] WebSocket连接已接受，客户端ID: {client_id}")
    
    client = None
    last_sent_text = ""

    # 定义结果回调函数
    async def send_result_to_client(text: str, sentence_end: bool = False):
        nonlocal last_sent_text
        try:
            if text != last_sent_text or sentence_end:
                message = json.dumps({
                    "type": "result",
                    "text": text,
                    "sentence_end": sentence_end,
                    "client_id": client_id
                })
                await websocket.send_text(message)
                last_sent_text = text if not sentence_end else ""
                logger.info(f"[阿里云ASR v2] 发送识别结果: {text} (sentence_end={sentence_end})")
        except Exception as e:
            logger.error(f"[阿里云ASR v2] 发送结果失败: {e}")

    try:
        while True:
            # logger.info(f"[阿里云ASR v2] 等待客户端 {client_id} 消息...")
            raw_msg = await websocket.receive_text()
            logger.info(f"[阿里云ASR v2] 收到客户端 {client_id} 消息: {raw_msg[:100]}...")
            
            try:
                msg_obj = json.loads(raw_msg)
            except json.JSONDecodeError as e:
                logger.error(f"[阿里云ASR v2] JSON解析失败: {e}")
                continue

            msg_type = msg_obj.get("type")
            # logger.info(f"[阿里云ASR v2] 消息类型: {msg_type}")
            
            if msg_type == "start":
                # 开始识别
                config = msg_obj.get("config", {})
                logger.info(f"[阿里云ASR v2] 开始识别，配置: {config}")
                
                try:
                    client = await asr_manager.create_client(
                        client_id,
                        model=config.get("model", "paraformer-realtime-v2"),
                        format=config.get("format", "pcm"),
                        sample_rate=config.get("sample_rate", 16000),
                        language_hints=config.get("language_hints", ["zh", "en"]),
                        on_result_callback=send_result_to_client
                    )
                    await client.connect()
                    
                    # 开始识别会话，参考官方demo
                    start_success = await client.start()
                    if not start_success:
                        raise Exception("识别会话启动失败")
                    
                    await websocket.send_text(json.dumps({
                        "type": "started",
                        "client_id": client_id
                    }))
                    logger.info(f"[阿里云ASR v2] 客户端 {client_id} 识别已开始")
                    
                except Exception as e:
                    logger.error(f"[阿里云ASR v2] 启动失败: {e}")
                    await websocket.send_text(json.dumps({
                        "type": "error",
                        "message": f"启动失败: {str(e)}",
                        "client_id": client_id
                    }))
                    
            elif msg_type == "audio":
                # 发送音频数据
                if client is None:
                    logger.warning(f"[阿里云ASR v2] 客户端 {client_id} 未初始化，请先调用start")
                    await websocket.send_text(json.dumps({
                        "type": "error",
                        "message": "请先调用start开始识别",
                        "client_id": client_id
                    }))
                    continue
                
                b64_data = msg_obj.get("data", "")
                # logger.info(f"[阿里云ASR v2] 收到音频数据，base64长度: {len(b64_data)}")
                try:
                    audio_bytes = base64.b64decode(b64_data)
                    # logger.info(f"[阿里云ASR v2] base64解码成功，音频数据长度: {len(audio_bytes)} 字节")
                    
                    # 检查音频数据是否为空或全零
                    if len(audio_bytes) == 0:
                        logger.warning(f"[阿里云ASR v2] 音频数据为空，跳过")
                        continue
                    
                    # 检查是否全为零（静音）
                    if all(b == 0 for b in audio_bytes):
                        logger.debug(f"[阿里云ASR v2] 检测到静音数据，长度: {len(audio_bytes)}")
                    else:
                        logger.info(f"[阿里云ASR v2] 检测到有效音频数据，长度: {len(audio_bytes)}")
                    
                    await client.send_audio(audio_bytes)
                    logger.info(f"[阿里云ASR v2] 音频数据已发送给阿里云ASR客户端，长度: {len(audio_bytes)}")
                    
                except Exception as e:
                    logger.error(f"[阿里云ASR v2] 处理音频数据失败: {e}")
                    logger.error(f"[阿里云ASR v2] base64数据前50字符: {b64_data[:50]}")
                    # 添加更详细的错误信息
                    import traceback
                    logger.error(f"[阿里云ASR v2] 错误详情: {traceback.format_exc()}")
                    
            elif msg_type == "stop":
                # 停止识别
                logger.info(f"[阿里云ASR v2] 客户端 {client_id} 请求停止识别")
                if client is not None:
                    # 先发送结束信号给阿里云
                    await client.stop()
                    # 然后移除客户端
                    await asr_manager.remove_client(client_id)
                    client = None
                    
                    await websocket.send_text(json.dumps({
                        "type": "stopped",
                        "client_id": client_id
                    }))
                    logger.info(f"[阿里云ASR v2] 客户端 {client_id} 识别已停止")
                    
            elif msg_type == "heartbeat":
                # 心跳包
                logger.info(f"[阿里云ASR v2] 收到心跳包，客户端ID: {client_id}")
                await websocket.send_text(json.dumps({
                    "type": "pong",
                    "client_id": client_id
                }))
                
            else:
                logger.warning(f"[阿里云ASR v2] 未知消息类型: {msg_type}")
                
    except WebSocketDisconnect:
        logger.info(f"[阿里云ASR v2] 客户端 {client_id} WebSocket已断开")
        if client_id in asr_manager.clients:
            await asr_manager.remove_client(client_id)
            
    except Exception as e:
        logger.error(f"[阿里云ASR v2] 处理WebSocket消息时发生错误: {e}")
        if client_id in asr_manager.clients:
            await asr_manager.remove_client(client_id) 